package io.sqrtqiezi.spark.streaming

import org.apache.spark.SparkConf
import org.apache.spark.streaming.{Seconds, State, StateSpec, StreamingContext}

object StateTracker2 {
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf()
      .setAppName("spark state tracker1")
      .setMaster("local[2]")

    val ssc = new StreamingContext(conf, Seconds(5))
    ssc.checkpoint("lagou-data/checkpoint2/")

    val wordDStream = ssc.socketTextStream("localhost", 9999)
      .flatMap(_.split("\\s+"))
      .map((_, 1))

    def mappingFunction(key: String, one: Option[Int], state: State[Int]): (String, Int) = {
      val sum = one.getOrElse(0) + state.getOption().getOrElse(0)
      state.update(sum)
      (key, sum)
    }

    val spec = StateSpec.function(mappingFunction _)
    val resultDStream = wordDStream.mapWithState(spec)
      .stateSnapshots()

    resultDStream.cache()

    resultDStream.print()

    resultDStream.repartition(1)
      .saveAsTextFiles("lagou-data/output2/")

    ssc.start()
    ssc.awaitTermination()
  }
}
